# -*- coding: utf-8 -*-
import torch
from torch import nn


class CustomNet(nn.Module):
    def __init__(self, num_classes=6):
        super().__init__()

        # 特征提取层（注意所有括号正确闭合）
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )  # 这里闭合features的Sequential

        # 自适应池化层（修正拼写错误）
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))

        # 分类器（确保所有括号闭合）
        self.classifier = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes)
        )  # 这里闭合classifier的Sequential

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


if __name__ == "__main__":
    # 测试模型结构
    model = CustomNet()
    test_input = torch.randn(1, 3, 224, 224)
    output = model(test_input)
    print(f"测试通过！输出形状: {output.shape}")
    print("模型结构验证：")
    print(model)